import akshare as ak
import pandas as pd
from sqlalchemy import create_engine, text

# MySQL 连接配置
engine = create_engine('mysql+pymysql://root:123456@localhost:3306/cn_stock_data')

# 示例批量季度
years = range(2025, 2026)
# quarters = ["0331", "0630", "0930", "1231"]
quarters = ["0630"]
report_dates = [f"{year}{q}" for year in years for q in quarters]

def process_profit_statement(report_date_str: str):
    print(f"开始处理 {report_date_str} ...")
    try:
        df = ak.stock_lrb_em(date=report_date_str)
        if df.empty:
            print(f"⚠️ {report_date_str} 没有数据，跳过")
            return

        if '序号' in df.columns:
            df = df.drop(columns=['序号'])

        # 数字列处理
        numeric_cols = [
            '净利润', '净利润同比', '营业总收入', '营业总收入同比',
            '营业总支出-营业支出', '营业总支出-销售费用', '营业总支出-管理费用',
            '营业总支出-财务费用', '营业总支出-营业总支出', '营业利润', '利润总额'
        ]
        for col in numeric_cols:
            if col in df.columns:
                df[col] = pd.to_numeric(df[col], errors='coerce')

        if '公告日期' in df.columns:
            df['公告日期'] = pd.to_datetime(df['公告日期'], errors='coerce').dt.date

        df['report_date'] = pd.to_datetime(report_date_str, format='%Y%m%d').date()

        df = df.rename(columns={
            '股票代码': 'stock_code',
            '股票简称': 'stock_name',
            '净利润': 'net_profit',
            '净利润同比': 'net_profit_yoy',
            '营业总收入': 'total_revenue',
            '营业总收入同比': 'total_revenue_yoy',
            '营业总支出-营业支出': 'total_expense',
            '营业总支出-销售费用': 'selling_expense',
            '营业总支出-管理费用': 'admin_expense',
            '营业总支出-财务费用': 'financial_expense',  # 表里字段 financial_expense
            '营业总支出-营业总支出': 'total_expense_sum', # 表里字段 total_expense_sum
            '营业利润': 'operating_profit',
            '利润总额': 'total_profit',
            '公告日期': 'announcement_date'
        })

        # 删除旧数据
        with engine.begin() as conn:
            conn.execute(
                text("DELETE FROM profit_statement WHERE report_date = :report_date"),
                {"report_date": df['report_date'].iloc[0]}
            )

        # 插入前只保留表中已有的列
        with engine.begin() as conn:
            table_cols = [row[0] for row in conn.execute(text("SHOW COLUMNS FROM profit_statement"))]
        df_to_insert = df[[c for c in df.columns if c in table_cols]]

        # 插入新数据
        df_to_insert.to_sql('profit_statement', con=engine, if_exists='append', index=False)
        print(f"✅ {report_date_str} 导入完成！ 共 {len(df)} 行")

    except Exception as e:
        print(f"❌ {report_date_str} 处理失败: {e}")

# 批量处理
for date in report_dates:
    process_profit_statement(date)

print("🎉 所有季度利润表导入完成！")
